#include <iostream>

#include "vortex/core/core.h"
#include "vortex/engine/classification_engine.h"
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

using namespace vortex;


class PaddleSwinInfer : public classificationEngine
{
public:
    PaddleSwinInfer(
        const std::string& engine_path, 
        const BlobInfo& input_info,
        const BlobInfo& output_info)
        : classificationEngine(engine_path, input_info, output_info) {}
    
    // custom preprocess functions
    virtual bool Preprocess(const std::string& image_path, int offset = 0) override
    {
        // default preprocessing
        cv::Mat image = cv::imread(image_path);
        uint32_t width = m_InputInfo.shape[3];
        uint32_t height = m_InputInfo.shape[2];
        cv::Mat temp;
        cv::resize(image, temp, cv::Size(width, height));

        uint32_t image_area = width * height;
        float* input_buffer = m_InputBlobs[0]->data_host;
        input_buffer = input_buffer + (image_area * 3) * offset;

        float* pBlue = input_buffer;
        float* pGreen = input_buffer + image_area;
        float* pRed = input_buffer + image_area * 2;

        unsigned char* pImage = temp.data;
        for (uint32_t i = 0; i < image_area; ++i)
        {
            pRed[i] = (pImage[3 * i + 0] / 255.0 - 0.5f) / 0.5f;
            pGreen[i] = (pImage[3 * i + 1] / 255.0 - 0.5f) / 0.5f;
            pBlue[i] = (pImage[3 * i + 2] / 255.0 - 0.5f) / 0.5f;
            // std::cout << pRed[i] << ", " << pGreen[i] << ", " << pBlue[i] << std::endl;
        }
        return true;
    }
};




int main()
{
    BlobInfo input_info = {"x", {16, 3, 224, 224}};
    BlobInfo output_info = {"softmax_25.tmp_0", {16, 21}};

#ifdef _WIN32
    std::string engine_path("D:/workspace/Tensorrt/vortex-onnx-trt/models/alexnet/alexnet.engine");
    std::string image_path("D:/workspace/Tensorrt/vortex-onnx-trt/data/samples/29bb3ece3180_11.jpg");
#else
    std::string engine_path("../../../models/paddleclas_swin/swin.engine");
    std::string image_path("../../../data/samples/buddism.jpg");
#endif

    PaddleSwinInfer model(engine_path, input_info, output_info);    
    std::vector<float> logits;
    
    model.Infer(image_path, logits);

    for (size_t i = 0; i < logits.size(); ++i)
        std::cout << logits[i] << std::endl;
    
    std::cout << "max label: \n";
    auto max = Argmax(logits);
    std::cout << max.first << ": " << max.second << std::endl;

    std::cout << "Main\n";
    return 0;
}
